# This file is part of Acacia+, a tool for synthesis of reactive systems using antichain-based techniques
# Copyright (C) 2011-2013 UMONS-ULB
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

from pygraph.classes.digraph import *
import subprocess
import os

from constants import *
from utils import *
  
#### Writes an mdp from transition system into a file (to give to PRISM)
def write_mdp_for_prism(ts_size, initial_states, states, label_values, sigma, sigma_I_size, nb_states, path, filename):
    mdp = ""

    nb_nondet_states = 0
    nb_add_states = 0 # number of additional states added (states reached from probabilistic states)
    nondet_states = dict()
    add_states = dict()
    add_states_map = dict() # used to merge equivalent add_states
    for key in sigma[P_I]:
        try:
            add_states_map[label_values[key]]
        except KeyError:
            add_states_map[label_values[key]] = dict()
    prob_dist = 1./sigma_I_size
        
    for state in states:
        nondet_states[state] = nb_nondet_states
        nb_nondet_states += 1
    
    if len(initial_states) == 1: # only one initial state
        init = nondet_states[initial_states[0]]
    else: # several initial states -> add a unique fake initial state
        init = nb_nondet_states
        nb_nondet_states += 1
      
    if len(initial_states) > 1:
        for cur_init in initial_states:
            mdp += "[ o_to_" + str(nondet_states[cur_init]) + " ] state = " + str(init) + " -> (state' = " + str(nondet_states[cur_init]) + ");\n"
        mdp += "\n"

    for cur_state in states:
        (player, transitions) = states[cur_state]
        if player == P_O: # deterministic, several actions might be available
            for cur_label in transitions:
                tos = transitions[cur_label]
                for to in tos:
                    mdp += "[ " + str(sigma[P_O][cur_label]) + " ] state = " + str(nondet_states[cur_state]) + " -> (state' = " + str(nondet_states[to]) + ");\n"
        else: # prob. dist., complete
            mdp += "\n"

            cur_add_states = []
            for cur_label in transitions:
                tos = transitions[cur_label]
                try:
                    cur_index_to = add_states_map[label_values[cur_label]][str(tos)]
                except KeyError:
                    cur_index_to = nb_nondet_states+nb_add_states
                    add_states_map[label_values[cur_label]][str(tos)] = cur_index_to                    
                    nb_add_states += 1
                for to in tos:
                    mdp += "[ " + str(sigma[P_I][cur_label]) + "_to_" + str(nondet_states[to]) + " ] state = " + str(cur_index_to) + " -> (state' = " + str(nondet_states[to]) + ");\n"
                cur_add_states.append(cur_index_to)
                
            mdp += "[ i ] state = " + str(nondet_states[cur_state]) + " -> "
            for i in range(sigma_I_size):
                mdp += str(prob_dist) + ": (state' = " + str(cur_add_states[i]) + ")"
                if i != sigma_I_size-1:
                    mdp += " + "
            mdp += ";\n"
            add_states[nondet_states[cur_state]] = cur_add_states
       
    mdp += "endmodule\n\n"
    
    mdp += "rewards \"cost\"\n"
    for label in sigma[P_O]:
        mdp += "[ " + str(sigma[P_O][label]) + " ] true : " + str(label_values[label]) + ";\n"
    mdp += "endrewards\n\n"
    
    mdp += "rewards \"reward\"\n"
    for label in sigma[P_O]:
        mdp += "[ " + str(sigma[P_O][label]) + " ] true : 1;\n"
    mdp += "endrewards\n"
   
    mdp = "mdp\n\nmodule strategies\nstate : [ 0 .. " + str(nb_nondet_states+nb_add_states-1) + " ] init "+str(init)+";\n\n" + mdp
              
    mdpfile = open(path + filename + ".nm", "w")                  
    mdpfile.write(mdp)
    mdpfile.close() 
        
    return (nondet_states, add_states, nb_nondet_states+nb_add_states, init)
    
#### Calls PRISM to compute best average case strategy from set of strategies
def call_prism(path, filename, nb_mdp_states, verbosity):
    start_time = start_time = os.times()[4]
    
    controled_print("Calling PRISM to solve MDP of winning strategies\n", [MINTEXT, ALLTEXT], verbosity)
    out = subprocess.Popen([PRISM_PATH+'prism/bin/prism', path+filename+".nm", '-pctl', 'R{"cost"}/{"reward"}min=? [S]', '-v'],stdout=subprocess.PIPE)
    (result, error) = out.communicate()
    
    prism_result_file = open(path + filename + ".prism", "w")              
    prism_result_file.write(result)
    prism_result_file.close() 

    result_lines = result.split("\n")
    for l in result_lines:
        if l.startswith("Value in the initial state"):
            expected_value = float(l.split(": ")[1])
    expected_value = -expected_value/2. # minus because prism objective is to minimize MP, divided by 2 because it is turn based (and no weight is associated to input signals) 
    
    result = "PRISM result:" + (result.split("PRISM")[0])#.split("Done")[1]
    
    controled_print(result+"\n", [ALLTEXT], verbosity)
    prism_time = os.times()[4]-start_time
    controled_print("PRISM time: %.2fs\n" % prism_time, [MINTEXT, ALLTEXT], verbosity)
    
    actions = [-1]*nb_mdp_states
    strategy = result.split("Action")[1:]

    for entry in strategy:
        s = entry.split("state =")
        r = s[0].split(", ")
        action = ""
        for a in r:
            if a.endswith("=1"):
                action = a.replace(".a=1", "") 
                break
        for i in range(1, len(s)):
            s[i] = s[i].replace(",", "")
            s[i] = s[i].replace(" ", "")
            s[i] = s[i].replace("\n", "")
            s[i] = s[i].replace("\t", "")
            actions[int(s[i])] = action

    return (actions, expected_value)


# Removes all actions that are not in the best average case strategy selected by PRISM (the optimal moves are stored in actions, the initial state is init_index)
def remove_non_optimal_actions(initial_states, states, sigma, starting_player, states_map, add_states, init_index, actions):
    states_map_inv = invert_dict(states_map)
    
    if len(initial_states) > 1: # Select initial state
        init_state = int(actions[init_index].split("_")[2])
        initial_states = [states_map_inv[init_state]]
    
    for state in states: # Select good transition for each state
        (player, out_tran) = states[state]
        if player == P_I: # Belongs to P_I -> remove non-determinism by selected the good transition
            i = 0
            for cur_sigma in sigma[P_I]:
                if len(out_tran[cur_sigma]) > 1:
                    good_action = actions[add_states[states_map[state]][i]].split("_")
                    out_tran[cur_sigma] = [states_map_inv[int(good_action[2])]] # The order in dictionnaries transitions and sigma[P_I] is the same (see their creation) -> it is ok
                    states[state] = (player, out_tran)
                i += 1
        else: # Belongs to P_O -> select good action between several choices
            good_action = actions[states_map[state]]
            if len(out_tran) > 1:
                key_to_pop = []
                for key in out_tran:
                    if sigma[P_O][key] != good_action:
                        key_to_pop.append(key)
                for key in key_to_pop:
                    out_tran.pop(key)
    
    return (initial_states, states)

# Creates a digraph from the strategy computed by PRISM and represented by states
def create_digraph_from_prism_strategy(initial_states, states, starting_player):
    g = digraph()
    g.__init__()
    
    for state_index in states :
        (sp, out_tran_sp) = states[state_index]
        state_sp = int(state_index)
        if sp == starting_player:
            try:
                g.add_node(state_sp)
                g.add_node_attribute(state_sp, ("label", str(len(g)-1)))
            except AdditionError:
                pass
    
            for label_sp in out_tran_sp:
                state_op = out_tran_sp[label_sp][0]
                (op, out_tran_op) = states[state_op]
                for label_op in out_tran_op:
                    state_to_sp = out_tran_op[label_op][0]
                    
                    try:
                        g.add_node(state_to_sp)
                        g.add_node_attribute(state_to_sp, ("label", str(len(g)-1)))
                    except AdditionError:
                        pass
                    
                    label = "((" + label_sp + ") U (" + label_op + "))"
                    if not g.has_edge((state_sp, state_to_sp)): # new edge from predecessor to new state
                        g.add_edge((state_sp, state_to_sp), 1, label)
                    else: # existing edge
                        new_label = g.edge_label((state_sp, state_to_sp)) + " || " + label
                        g.set_edge_label((state_sp, state_to_sp), new_label)  
                        
    g.add_node_attribute(initial_states[0], ('initial', True))
    
    return g